import time

import rtde_control
import rtde_receive
import sys
import threading
import numpy as np
from compensation import GravityCompensation
from Hardware.KunWei_ftsensor import ftsensor

ip='192.168.68.250'

rtde_r = rtde_receive.RTDEReceiveInterface(ip)
rtde_c = rtde_control.RTDEControlInterface(ip)
compensation = GravityCompensation()
ft = ftsensor(False)

# force = rtde_r.getFtRawWrench()
time.sleep(1)
force = ft.GetForce()
q = rtde_r.getActualQ()
rotation_vector = rtde_c.getForwardKinematics(q,[0,0,0,0,0,0])
compensation.FtZero(force,rotation_vector[3:6])

rtde_c.teachMode()
while True:
    force = ft.GetForce()
    q = rtde_r.getActualQ()
    rotation_vector = rtde_c.getForwardKinematics(q, [0, 0, 0, 0, 0, 0])
    force = compensation.GetForceTorque(force, rotation_vector[3:6])
    sys.stdout.write("\r" + str(force))
    sys.stdout.flush()